# react_mini_deepseek.py
# 最小可用 ReAct：DeepSeek + JSON 动作 + 两个工具（calculator / time_now）
# 依赖：pip install openai python-dotenv

import os
import json
import ast
import datetime
import textwrap
from dataclasses import dataclass
from typing import Callable, Dict, Any, Optional, Tuple
from zoneinfo import ZoneInfo

from dotenv import load_dotenv
load_dotenv()

# ==== 配置：DeepSeek 的 OpenAI 兼容调用 ====
try:
    from openai import OpenAI
except Exception as e:
    raise SystemExit("请先安装依赖：pip install openai") from e

API_KEY = os.getenv("DEEPSEEK_API_KEY")
BASE_URL = os.getenv("DEEPSEEK_BASE_URL")  # 按 DeepSeek 官方文档设置
MODEL = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")

if not API_KEY or not BASE_URL:
    raise SystemExit(
        "请设置环境变量 DEEPSEEK_API_KEY 与 DEEPSEEK_BASE_URL（详见 DeepSeek 文档）。"
    )

client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

# ==== 工具定义 ====

@dataclass
class Tool:
    name: str
    description: str
    schema: Dict[str, Any]  # 简化：仅用于提示，不做强校验
    func: Callable[[Dict[str, Any]], str]

def safe_eval_expr(expr: str) -> float:
    """
    安全算式求值：仅允许 + - * / ** 与括号、数值常量。
    禁止变量名、函数调用，以及非数值常量（如字符串、布尔、None）。
    """
    allowed_nodes = {
        ast.Expression, ast.BinOp, ast.UnaryOp, ast.Load,
        ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.USub, ast.UAdd,
        ast.Constant,  # 仅允许数值型常量，见下面的类型检查
    }

    tree = ast.parse(expr, mode='eval')

    for node in ast.walk(tree):
        if isinstance(node, ast.Call):
            raise ValueError("不允许函数调用")
        if isinstance(node, ast.Name):
            raise ValueError("不允许变量名")
        if type(node) not in allowed_nodes:
            raise ValueError(f"不允许的语法: {type(node).__name__}")
        # 关键：仅放行数值常量（int/float）
        if isinstance(node, ast.Constant) and not isinstance(node.value, (int, float)):
            raise ValueError("只允许数值常量（int/float）")

    return eval(compile(tree, "<expr>", "eval"), {"__builtins__": {}}, {})


def tool_calculator(params: Dict[str, Any]) -> str:
    expr = str(params.get("expression", "")).strip()
    if not expr:
        return "错误：缺少 expression"
    try:
        val = safe_eval_expr(expr)
        return f"结果={val}"
    except Exception as e:
        return f"错误：无法计算（{e}）"

def tool_time_now(params: Dict[str, Any]) -> str:
    zone = params.get("zone")
    try:
        if zone:
            now = datetime.datetime.now(tz=ZoneInfo(zone))
        else:
            now = datetime.datetime.now()  # 本地时区
        return now.isoformat()
    except Exception as e:
        return f"错误：无效时区（{e}）"

TOOLS: Dict[str, Tool] = {
    "calculator": Tool(
        name="calculator",
        description="安全计算算式，支持 + - * / ** 与括号；参数：expression (string)。",
        schema={"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"]},
        func=tool_calculator
    ),
    "time_now": Tool(
        name="time_now",
        description="返回当前时间（ISO 8601 字符串）；可选参数：zone (IANA 时区，如 'Asia/Shanghai')。",
        schema={"type": "object", "properties": {"zone": {"type": "string"}}},
        func=tool_time_now
    ),
}

# ==== Prompt 设计（系统提示词中文） ====

SYSTEM_PROMPT = """\
你是一个遵循 ReAct 模式的助理。必须先逐步思考（thought），再选择行动（action），并通过工具完成任务。
你必须只输出一个 JSON 对象，且不得包含任何额外文本或 Markdown。
JSON 结构如下：
{"thought": "...", "action": "<工具名或 final>", "action_input": {... 或 "<最终答案文本>"}}

规则：
- 若需要使用工具：将 "action" 设为该工具名，将 "action_input" 设为该工具的 JSON 参数对象。
- 若可以直接给出最终答案：将 "action" 设为 "final"，并把最终答案文本放入 "action_input"。
- 只允许这三个键：thought、action、action_input；严禁出现其它键。
- 严禁输出解释性文字、前后缀、或代码块标记（例如 ```）。
"""

def build_tools_block(tools: Dict[str, Tool]) -> str:
    lines = ["可用工具列表："]
    for t in tools.values():
        lines.append(f"- {t.name}: {t.description}")
        schema = json.dumps(t.schema, ensure_ascii=False)
        lines.append(f"  schema={schema}")
    return "\n".join(lines)

def build_user_prompt(task: str, scratchpad: str) -> str:
    return textwrap.dedent(f"""\
    任务：{task}

    {build_tools_block(TOOLS)}

    现在请只输出 JSON（不要解释，不要 markdown），字段为 thought/action/action_input。
    历史轨迹（供你参考）：
    {scratchpad if scratchpad.strip() else "（无）"}
    """).strip()

# ==== LLM 调用与解析 ====

def llm_json_decision(messages) -> Dict[str, Any]:
    """
    调用 DeepSeek（OpenAI 兼容），解析出 JSON 对象；
    若第一次未返回合规 JSON，进行一次自我修复重试。
    """
    resp = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.2,
        top_p=0.9,
    )
    content = resp.choices[0].message.content.strip()

    def try_parse(txt: str) -> Optional[Dict[str, Any]]:
        try:
            left = txt.find("{")
            right = txt.rfind("}")
            if left != -1 and right != -1 and right > left:
                txt = txt[left:right+1]
            return json.loads(txt)
        except Exception:
            return None

    data = try_parse(content)
    if data is None:
        repair_messages = messages + [
            {"role": "assistant", "content": content},
            {"role": "user", "content": "你没有按要求输出可解析 JSON。请严格只输出 JSON（thought/action/action_input）。"}
        ]
        resp2 = client.chat.completions.create(
            model=MODEL,
            messages=repair_messages,
            temperature=0.0
        )
        content2 = resp2.choices[0].message.content.strip()
        data = try_parse(content2)

    if not isinstance(data, dict) or not all(k in data for k in ("thought", "action", "action_input")):
        preview = content[:200].replace("\n", " ")
        raise ValueError(f"模型未返回正确 JSON：{preview}...")

    return data

# ==== ReAct 主循环 ====

def react_loop(task: str, max_steps: int = 10, verbose: bool = True) -> Tuple[str, str]:
    """
    执行 ReAct 闭环，返回 (final_answer, trajectory_text)
    """
    scratchpad_steps = []  # 展示轨迹：Thought/Action/Observation
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
    ]

    for step in range(1, max_steps + 1):
        scratchpad_text = "\n".join(scratchpad_steps)
        user_prompt = build_user_prompt(task, scratchpad_text)
        messages_step = messages + [{"role": "user", "content": user_prompt}]

        decision = llm_json_decision(messages_step)
        thought = str(decision.get("thought", "")).strip()
        action = str(decision.get("action", "")).strip()
        action_input = decision.get("action_input")

        if verbose:
            print(f"\n[Step {step}] Thought: {thought}")
            print(f"[Step {step}] Action:  {action}")
            print(f"[Step {step}] Action Input: {action_input}")

        if action == "final":
            final_answer = str(action_input)
            scratchpad_steps.append(f"Thought: {thought}\nFinal: {final_answer}")
            return final_answer, "\n".join(scratchpad_steps)

        tool = TOOLS.get(action)
        if not tool:
            obs = f"错误：未知工具 {action}"
        else:
            if not isinstance(action_input, dict):
                obs = "错误：action_input 不是 JSON 对象"
            else:
                try:
                    obs = tool.func(action_input)
                except Exception as e:
                    obs = f"错误：工具执行异常（{e}）"

        scratchpad_steps.append(
            f"Thought: {thought}\nAction: {action}\nAction Input: {json.dumps(action_input, ensure_ascii=False)}\nObservation: {obs}"
        )

    final = "达到最大步数仍未完成，请缩小问题或增加步数。"
    return final, "\n".join(scratchpad_steps)

# ==== 运行示例 ====

if __name__ == "__main__":
    demo_task = "先计算 3.5 的平方，再加上我所在时区（Asia/Shanghai）的当前小时数，然后给出最终结果。"
    answer, trace = react_loop(demo_task, max_steps=5, verbose=True)

    print("\n====== 最终答案 ======")
    print(answer)
    print("\n====== 轨迹回放 ======")
    print(trace)
